from tqdm import tqdm
import sys
# from pys.paraformer 
import utils_file
from gxl_wv_mos import get_wvmos
# from project.data_handle_pipline import get_scp_file
wavs_scp = sys.argv[1]
out_scp = sys.argv[2]
gpuid = sys.argv[3]
model = get_wvmos(cuda=True, device=f"cuda:{gpuid}")


def get_mos_score_for_single_wav(wav_path):
    return model.calculate_one(wav_path)


def get_mvmos_for_scp_file(wav_scp_file_path: str, target_scp_file_path: str):
    utils_file.makedir_for_file(target_scp_file_path)
    dic_data = utils_file.load_dic_from_pathscp(wav_scp_file_path)
    res_dic = {}
    num = 0
    print('开始得到得到wav.scp文件对应音频的mvmos, 共有{}条数据'.format(len(dic_data)))
    for k, v in tqdm(dic_data.items(), total=len(dic_data)):
        output = get_mos_score_for_single_wav(v)
        res_dic[k] = output
    utils_file.write_dic_to_scp(res_dic, target_scp_file_path)


if __name__ == '__main__':
    # print(get_mos_score_for_single_wav("./output/wav/dev/BAC009S0724W0121.wav"))  # infer MOS score for one audio
    #  print(get_mos_score_for_single_wav("/home/work_nfs7/xlgeng/code_runner_gxl/ai/output/gxl.wav"))  # infer MOS score for one audio
    #  print(get_mos_score_for_single_wav("/home/work_nfs7/xlgeng/code_runner_gxl/ai/output/BAC009S0724W0126.wav"))  # infer MOS score for one audio
    # get_scp_file("/home/work_nfs6/lhma/wash_data/sample1000/res/muise_denoise","./output/wav.scp")

    get_mvmos_for_scp_file(wavs_scp, out_scp)